#!/usr/bin/env python3
"""
图片生成任务 - 异步处理
处理图片生成和模板图片融合
"""

from celery import current_app, current_task
from app.celery_app import celery_app
from app.services.ai_manager import ai_service_manager
from app.services.ai_base import GenerationRequest, AIServiceType, AIServiceProvider
import logging
import asyncio
from typing import Dict, Any

logger = logging.getLogger(__name__)

@celery_app.task(bind=True, name='generate_image_async')
def generate_image_async(self, request_data: Dict[Any, Any], user_id: int, params: Dict[str, Any]):
    """
    图片生成异步任务
    支持文本生图和模板图片融合
    """
    task_id = self.request.id
    logger.info(f"开始执行图片生成任务 {task_id}")
    
    try:
        # 更新任务状态：开始处理
        self.update_state(
            state='PROGRESS',
            meta={
                'progress': 10,
                'message': '开始图片生成...',
                'step': 'image_generation'
            }
        )
        
        # 判断生成类型
        is_template_mode = bool(
            params.get("template_image_url") or params.get("template_image_base64")
        )
        
        if is_template_mode:
            logger.info(f"任务 {task_id} - 执行以图生图")
            generation_type = "image_to_image"  # 即梦AI的以图生图
        else:
            logger.info(f"任务 {task_id} - 执行文本生图")
            generation_type = "text_to_image"
        
        # 构建图片生成请求
        image_request = GenerationRequest(
            prompt=params["prompt"],
            service_type=AIServiceType.IMAGE_GENERATION,
            params={
                "template_image_url": params.get("template_image_url"),
                "user_image_url": params.get("user_image_url"),
                "template_image_base64": params.get("template_image_base64"),
                "user_image_base64": params.get("user_image_base64"),
                "aspect_ratio": params["aspect_ratio"],
                "style": params["style"],
                "fusion_strength": params.get("fusion_strength", 0.6),
                "seed": params["seed"],
                "generation_type": generation_type
            }
        )
        
        # 更新进度
        self.update_state(
            state='PROGRESS',
            meta={
                'progress': 30,
                'message': '正在生成图片...',
                'step': 'processing'
            }
        )
        
        # 执行图片生成
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        
        try:
            # 获取即梦AI服务
            jimeng_service = ai_service_manager.get_service(AIServiceProvider.JIMENG_AI)
            if not jimeng_service:
                raise Exception("即梦AI服务不可用")
            
            if is_template_mode:
                result = loop.run_until_complete(
                    jimeng_service.generate_image_fusion(image_request)
                )
            else:
                result = loop.run_until_complete(
                    jimeng_service.generate_image(image_request)
                )
            
            if not result.success:
                raise Exception(f"图片生成失败: {result.error_message}")
                
            image_url = result.result_url
            cost = result.cost or 0.1
            
            logger.info(f"任务 {task_id} - 图片生成完成: {image_url}")
            
            # 返回最终结果
            final_result = {
                'success': True,
                'image_url': image_url,
                'cost': cost,
                'message': '图片生成完成',
                'generation_type': generation_type,
                'aspect_ratio': params['aspect_ratio']
            }
            
            logger.info(f"任务 {task_id} - 图片生成任务执行成功")
            return final_result
            
        finally:
            loop.close()
    
    except Exception as e:
        logger.error(f"任务 {task_id} - 执行失败: {str(e)}")
        
        failure_meta = {
            'exc_type': type(e).__name__,
            'exc_message': str(e),
            'message': f'图片生成失败: {str(e)}',
            'success': False
        }

        self.update_state(
            state='FAILURE',
            meta=failure_meta
        )
        
        return {
            'success': False,
            'error': str(e),
            'message': f'图片生成失败: {str(e)}'
        }